"""
重排模型调用模块 - 负责调用不同的重排模型API并保存结果
"""
import csv
import json
import os
import requests
from typing import List, Dict, Any, Optional, Tuple
import pandas as pd

# 导入API客户端模块
from src.api_client import call_bge_rerank_api, call_qwen_rerank_api

def call_bge_rerank(question: str, passages: List[str]) -> Dict[str, Any]:
    """
    调用BGE重排模型API
    
    参数:
        question: 问题文本
        passages: 知识块列表
        
    返回:
        包含scores和rankings的字典
    """
    print(f"使用BGE重排模型为问题 '{question}' 重排 {len(passages)} 个知识块")
    
    # 调用实际的BGE重排API
    scores = call_bge_rerank_api(question, passages, verify_ssl=False, max_retries=3)
    
    # 根据分数生成排名（索引从0开始）
    # 创建索引和分数的配对，然后按分数降序排序
    indexed_scores = [(i, score) for i, score in enumerate(scores)]
    indexed_scores.sort(key=lambda x: x[1], reverse=True)
    
    # 提取排名列表（按重排分数从高到低的索引顺序）
    rankings = [index for index, score in indexed_scores]
    
    return {
        "scores": scores,
        "rankings": rankings
    }

def call_qwen_rerank(question: str, passages: List[str]) -> Dict[str, Any]:
    """
    调用Qwen重排模型API
    
    参数:
        question: 问题文本
        passages: 知识块列表
        
    返回:
        包含scores和rankings的字典
    """
    print(f"使用Qwen重排模型为问题 '{question}' 重排 {len(passages)} 个知识块")
    
    # 调用实际的Qwen重排API
    scores = call_qwen_rerank_api(question, passages, verify_ssl=False, max_retries=3)
    
    # 根据分数生成排名（索引从0开始）
    # 创建索引和分数的配对，然后按分数降序排序
    indexed_scores = [(i, score) for i, score in enumerate(scores)]
    indexed_scores.sort(key=lambda x: x[1], reverse=True)
    
    # 提取排名列表（按重排分数从高到低的索引顺序）
    rankings = [index for index, score in indexed_scores]
    
    return {
        "scores": scores,
        "rankings": rankings
    }

def read_evaluation_results(csv_path: str) -> Dict[str, Dict[str, Any]]:
    """
    从评估结果CSV文件中读取数据
    
    参数:
        csv_path: CSV文件路径
        
    返回:
        问题和知识块的嵌套字典
    """
    data = {}
    df = pd.read_csv(csv_path)
    
    for _, row in df.iterrows():
        question = row['问题']
        chunk = row['知识块']
        score = float(row['评分'])
        reason = row['理由']
        
        if question not in data:
            data[question] = {
                "passages": [],
                "gold_scores": [],
                "reasons": []
            }
        
        data[question]["passages"].append(chunk)
        data[question]["gold_scores"].append(score)
        data[question]["reasons"].append(reason)
    
    print(f"从 {csv_path} 读取了 {len(data)} 个问题的评估结果")
    return data

def save_rerank_results(data: Dict[str, Dict[str, Any]], 
                       bge_results: Dict[str, Dict[str, Any]], 
                       qwen_results: Dict[str, Dict[str, Any]], 
                       output_csv: str):
    """
    将重排结果保存为CSV文件
    
    参数:
        data: 原始评估数据
        bge_results: BGE重排结果
        qwen_results: Qwen重排结果
        output_csv: 输出CSV文件路径
    """
    rows = []
    
    for question_id, question_data in data.items():
        passages = question_data["passages"]
        gold_scores = question_data["gold_scores"]
        
        bge_result = bge_results.get(question_id, {"scores": [], "rankings": []})
        qwen_result = qwen_results.get(question_id, {"scores": [], "rankings": []})
        
        for i, passage in enumerate(passages):
            row = {
                "问题": question_id,  # 保持与评估结果CSV一致的列名
                "知识块索引": i,
                "知识块": passage,
                "黄金标准评分": gold_scores[i] if i < len(gold_scores) else 0,
                "BGE重排分数": bge_result["scores"][i] if i < len(bge_result["scores"]) else 0,
                "BGE排名": bge_result["rankings"].index(i) if i in bge_result["rankings"] else -1,
                "Qwen重排分数": qwen_result["scores"][i] if i < len(qwen_result["scores"]) else 0,
                "Qwen排名": qwen_result["rankings"].index(i) if i in qwen_result["rankings"] else -1
            }
            rows.append(row)
    
    # 保存为CSV（Excel友好编码）
    df = pd.DataFrame(rows)
    df.to_csv(output_csv, index=False, encoding='utf-8-sig')
    
    print(f"重排结果已保存到 {output_csv}")

def run_reranking_pipeline(input_csv: str, output_csv: str):
    """
    运行完整的重排流程
    
    参数:
        input_csv: 输入CSV文件路径（评估结果）
        output_csv: 输出CSV文件路径（重排结果）
    """
    # 1. 读取评估结果
    data = read_evaluation_results(input_csv)
    
    # 2. 对每个问题进行重排
    bge_results = {}
    qwen_results = {}
    
    for question_id, question_data in data.items():
        passages = question_data["passages"]
        
        # 2.1 调用BGE重排
        bge_result = call_bge_rerank(question_id, passages)
        bge_results[question_id] = bge_result
        
        # 2.2 调用Qwen重排
        qwen_result = call_qwen_rerank(question_id, passages)
        qwen_results[question_id] = qwen_result
    
    # 3. 保存重排结果
    save_rerank_results(data, bge_results, qwen_results, output_csv)
    
    print(f"重排流程完成，结果已保存到 {output_csv}")

if __name__ == "__main__":
    # 示例用法
    input_csv = "data/processed/evaluation_results.csv"
    output_csv = "data/processed/reranking_results.csv"
    
    run_reranking_pipeline(input_csv, output_csv) 